from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.distance import pdist, squareform
from scipy.stats import pearsonr
from scipy.interpolate import Rbf


energy_level_trial1 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    0, # 6
    3, # 7
    0, # 8
    3, # 9
    3, # 10
    0, # 11
    0, # 12
    0, # 13
    2, # 14
    2, # 15
    0, # 16
    1, # 17
    1, # 18
    0, # 19
    0.5, # 20
    0.5, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial2 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    0, # 6
    3, # 7
    3, # 8
    3, # 9
    3, # 10
    2, # 11
    0, # 12
    1, # 13
    1, # 14
    2, # 15
    2, # 16
    1, # 17
    0, # 18
    0, # 19
    0, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial3 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    0, # 6
    0, # 7
    3, # 8
    3, # 9
    0, # 10
    0, # 11
    0, # 12
    3, # 13
    3, # 14
    0, # 15
    0, # 16
    0, # 17
    0, # 18
    0, # 19
    0, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial4 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    0, # 6
    0, # 7
    0, # 8
    0, # 9
    0, # 10
    0, # 11
    0, # 12
    0, # 13
    0, # 14
    0, # 15
    0, # 16
    0, # 17
    0, # 18
    0, # 19
    0, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial5 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    0, # 6
    0, # 7
    3, # 8
    3, # 9
    3, # 10
    0, # 11
    0, # 12
    2, # 13
    2, # 14
    2, # 15
    1, # 16
    1, # 17
    0, # 18
    0, # 19
    0, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial6 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    0, # 6
    0, # 7
    3, # 8
    3, # 9
    3, # 10
    3, # 11
    0, # 12
    0, # 13
    0, # 14
    1, # 15
    0, # 16
    0, # 17
    0, # 18
    0, # 19
    0, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial7 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    0, # 6
    3, # 7
    3, # 8
    3, # 9
    3, # 10
    2, # 11
    0, # 12
    0, # 13
    0, # 14
    2, # 15
    2, # 16
    1, # 17
    1, # 18
    1, # 19
    0, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial8 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    0, # 6
    0, # 7
    3, # 8
    3, # 9
    3, # 10
    0, # 11
    0, # 12
    2, # 13
    2, # 14
    2, # 15
    2, # 16
    1, # 17
    1, # 18
    0, # 19
    0, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial9 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    3, # 6
    0, # 7
    3, # 8
    3, # 9
    3, # 10
    3, # 11
    0, # 12
    2, # 13
    2, # 14
    2, # 15
    2, # 16
    0, # 17
    0, # 18
    0, # 19
    0, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_level_trial10 = np.array([
    0, # 0
    0, # 1
    0, # 2
    0, # 3
    0, # 4
    0, # 5
    3, # 6
    3, # 7
    3, # 8
    3, # 9
    3, # 10
    3, # 11
    0, # 12
    1, # 13
    1, # 14
    2, # 15
    2, # 16
    2, # 17
    0, # 18
    1, # 19
    1, # 20
    0, # 21
    0, # 22
    0, # 23
    ])

energy_levels = [energy_level_trial1, energy_level_trial2, energy_level_trial3, energy_level_trial4, energy_level_trial5, energy_level_trial6, energy_level_trial7, energy_level_trial8, energy_level_trial9, energy_level_trial10]

energy_levels_mean = np.mean(energy_levels, axis=0)

# 使用 径向基函数 拟合 energy_levels_mean
x = np.arange(24)
rbf = Rbf(x, energy_levels_mean, function='gaussian')
y_rbf = rbf(x)

# 将每一个 energy level 绘制到同一个plot中，用折线图表示

fig = plt.figure()
ax = fig.add_subplot(111)
x = np.arange(24)
ax.plot(x, energy_level_trial1, label='trial1')
ax.plot(x, energy_level_trial2, label='trial2')
ax.plot(x, energy_level_trial3, label='trial3')
ax.plot(x, energy_level_trial4, label='trial4')
ax.plot(x, energy_level_trial5, label='trial5')
ax.plot(x, energy_level_trial6, label='trial6')
ax.plot(x, energy_level_trial7, label='trial7')
ax.plot(x, energy_level_trial8, label='trial8')
ax.plot(x, energy_level_trial9, label='trial9')
ax.plot(x, energy_level_trial10, label='trial10')
plt.legend(loc='upper right')
plt.show()

fig = plt.figure()
ax = fig.add_subplot(111)
x = np.arange(24)

for energy_level in energy_levels:
    ax.bar(x, energy_level, color='gray', alpha=0.1)

ax.plot(y_rbf, label='RBF', linewidth=2.0)

plt.legend(loc='upper right')
plt.show()
